from PIL import Image, ImageFilter, ImageOps
import numpy as np
import io
import requests
from dotenv import load_dotenv
import os

load_dotenv()

MIDJOURNEY_API_KEY = os.getenv('MIDJOURNEY_API_KEY')


def convert_to_comic(image_path, style="default"):
    """将上传的图像转换为漫画风格"""
    try:
        # 根据风格选择处理方式
        if style == "black-white":
            return classic_comic_effect(image_path)
        elif style == "watercolor":
            return watercolor_effect(image_path)
        elif style == "pop-art":
            return pop_art_effect(image_path)
        else:  # 默认使用Midjourney风格转换
            return midjourney_style_transfer(image_path, style)
    except Exception as e:
        print(f"Error converting image: {e}")
        return None


def midjourney_style_transfer(image_path, style):
    """使用Midjourney API进行风格转换"""
    try:
        headers = {
            "Authorization": f"Bearer {MIDJOURNEY_API_KEY}",
            "Content-Type": "application/json"
        }

        # 读取图像并转换为base64
        with open(image_path, "rb") as img_file:
            img_base64 = base64.b64encode(img_file.read()).decode('utf-8')

        style_prompts = {
            "default": "comic book style, vibrant colors, bold outlines",
            "shonen": "shonen manga style, dynamic action, speed lines",
            "shoujo": "shoujo manga style, delicate features, flowery backgrounds",
            "chibi": "chibi style, super deformed, cute exaggeration"
        }

        prompt = style_prompts.get(style, style_prompts["default"])

        payload = {
            "image": img_base64,
            "prompt": f"{prompt} --v 5",
            "process": "style_transfer"
        }

        response = requests.post(
            "https://api.midjourney.com/v1/transform",
            headers=headers,
            json=payload
        )

        if response.status_code == 200:
            return io.BytesIO(response.content)
        else:
            print(f"Midjourney API error: {response.text}")
            return classic_comic_effect(image_path)  # 回退到本地处理
    except Exception as e:
        print(f"Error in Midjourney style transfer: {e}")
        return classic_comic_effect(image_path)


def classic_comic_effect(image_path):
    """经典黑白漫画效果"""
    img = Image.open(image_path)

    # 转换为灰度
    img = img.convert('L')

    # 边缘检测
    edges = img.filter(ImageFilter.FIND_EDGES)
    edges = ImageOps.invert(edges)

    # 半色调效果
    img = img.resize((img.width // 2, img.height // 2), Image.BILINEAR)
    img = img.resize((img.width * 2, img.height * 2), Image.NEAREST)

    # 合并效果
    result = Image.blend(img, edges, 0.7)

    output = io.BytesIO()
    result.save(output, format='PNG')
    output.seek(0)

    return output


def watercolor_effect(image_path):
    """水彩画效果"""
    img = Image.open(image_path)

    # 轻微模糊模拟水彩笔触
    img = img.filter(ImageFilter.SMOOTH_MORE)

    # 增强颜色
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.5)

    output = io.BytesIO()
    img.save(output, format='PNG')
    output.seek(0)

    return output


def pop_art_effect(image_path):
    """波普艺术效果"""
    img = Image.open(image_path)

    # 减少颜色数量
    img = img.convert('P', palette=Image.ADAPTIVE, colors=8)
    img = img.convert('RGB')

    # 增强对比度
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(2.0)

    output = io.BytesIO()
    img.save(output, format='PNG')
    output.seek(0)

    return output